import akshare as ak
import pandas as pd
import finplot as fplt
import numpy as np
from PIL import Image
import time
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
import mplfinance as mpf
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib import gridspec
# 设置字体，确保中文正常显示
plt.rcParams["font.family"] = ["simhei"]

# 导入random模块用于随机化
import random

def get_stock_data(symbol="600036", period="daily", start_date="20250101", end_date="20250731", adjust="qfq"):
    """
    获取股票数据

    Parameters:
    symbol (str): 股票代码
    period (str): 时间周期 ("daily", "weekly", "monthly")
    start_date (str): 开始日期 (YYYYMMDD)
    end_date (str): 结束日期 (YYYYMMDD)
    adjust (str): 复权方式 ("qfq": 前复权, "hfq": 后复权, "": 不复权)

    Returns:
    pd.DataFrame: 股票数据
    """
    df = ak.stock_zh_a_hist(
        symbol=symbol,
        period=period,
        start_date=start_date,
        end_date=end_date,
        adjust=adjust
    )

    # 将中文列名转换为英文列名
    column_mapping = {
        '日期': 'date',
        '股票代码': 'symbol',
        '开盘': 'open',
        '收盘': 'close',
        '最高': 'high',
        '最低': 'low',
        '成交量': 'volume',
        '成交额': 'amount',
        '振幅': 'amplitude',
        '涨跌幅': 'change_rate',
        '涨跌额': 'change_price',
        '换手率': 'turnover_rate'
    }

    # 重命名列
    df = df.rename(columns=column_mapping)
    # 将date 转为时间

    df['timestamp'] = pd.to_datetime(df['date'])
    # 将timestamp 设置成时间索引
    df = df.set_index('timestamp')
    # df = df.reset_index()
    # df = df.rename(columns={"index": "timestamp"})

    return df


def calculate_ene(df, N=10, M1=11, M2=9):
    """
    计算ENE轨道线指标
    ENE轨道线是基于移动平均线和一定比例的上下轨道线构成的技术分析指标
    
    Parameters:
    df (pd.DataFrame): 股票数据
    N (int): 移动平均周期
    M1 (int): 上轨道百分比
    M2 (int): 下轨道百分比
    
    Returns:
    pd.DataFrame: 添加了ENE指标的数据
    """
    df = df.copy()
    
    # 计算移动平均线
    df['MA'] = df['close'].rolling(window=N).mean()
    
    # 计算上轨道线 (UPPER)
    df['ENE_UPPER'] = (1 + M1 / 100) * df['MA']
    
    # 计算下轨道线 (LOWER)
    df['ENE_LOWER'] = (1 - M2 / 100) * df['MA']
    
    # 计算ENE中线 (ENE线)
    df['ENE'] = (df['ENE_UPPER'] + df['ENE_LOWER']) / 2
    
    return df

def calculate_moving_averages(df, windows=[5, 10, 20]):
    """
    计算价格均线
    
    Parameters:
    df (pd.DataFrame): 股票数据
    windows (list): 均线周期列表
    
    Returns:
    pd.DataFrame: 添加了均线的数据
    """
    df = df.copy()
    
    # 计算价格均线
    for window in windows:
        df[f'ma_price_{window}'] = df['close'].rolling(window=window).mean()
        
    return df

def calculate_pressure_lines(df, period=20):
    """
    计算近期压力线（基于近期最高价和最低价）
    
    Parameters:
    df (pd.DataFrame): 股票数据
    period (int): 计算周期
    
    Returns:
    tuple: (近期最高值, 近期最低值)
    """
    df = df.copy()
    
    # 计算近期最高价和最低价
    recent_high = df['high'].tail(period).max()
    recent_low = df['low'].tail(period).min()
    
    return recent_high, recent_low

def plot_stock_chart(df_data, symbol="600036"):
    """
    使用finplot绘制股票图表，包含K线、均线、ENE轨道线和压力线
    
    Parameters:
    df_data (pd.DataFrame): 股票数据
    symbol (str): 股票代码
    """
    df = df_data.copy()
    # 只取最近100天的数据
    df = df.tail(100)
    
    # 计算压力线
    pressure_high, pressure_low = calculate_pressure_lines(df)
    
    # 创建图表
    ax1, ax2, ax3 = fplt.create_plot(f'{symbol} 股票分析图表', rows=3, init_zoom_periods=100)
    candle_style = {'up_color': '#ff0000', 'down_color': '#00ff00',
                    'edgecolor': '#000000', 'wickcolor': '#000000'}
    # 绘制K线图
    fplt.candlestick_ochl(df[['open', 'close', 'high', 'low']], ax=ax1,style=candle_style)
    
    # 绘制价格均线
    fplt.plot(df['ma_price_5'], ax=ax1, width=5, color='#0000ff', legend='MA5')
    fplt.plot(df['ma_price_10'], ax=ax1,width=4, color='#ffa500', legend='MA10')
    fplt.plot(df['ma_price_20'], ax=ax1, width=3,color='#800080', legend='MA20')
    
    # 绘制ENE轨道线
    fplt.plot(df['ENE_UPPER'], ax=ax1, width=6,color='#ff0000', legend='ENE_UPPER')
    fplt.plot(df['ENE_LOWER'], ax=ax1, width=6,color='#008000', legend='ENE_LOWER')
    fplt.plot(df['ENE'], ax=ax1, width=6,color='#000000', legend='ENE')
    
    # 绘制压力线
    fplt.plot(pd.Series([pressure_high] * len(df)), ax=ax1, color='#ff0000', style='--', legend=f'压力线: {pressure_high:.2f}')
    fplt.plot(pd.Series([pressure_low] * len(df)), ax=ax1, color='#008000', style='--', legend=f'支撑线: {pressure_low:.2f}')
    
    # 绘制成交量
    fplt.volume_ocv(df[['open', 'close', 'volume']], ax=ax2)
    
    # 计算并绘制量均线
    df['ma_volume_5'] = df['volume'].rolling(window=5).mean()
    df['ma_volume_10'] = df['volume'].rolling(window=10).mean()
    fplt.plot(df['ma_volume_5'], ax=ax2,  width=5,color='#ffa500', legend='Volume MA5')
    fplt.plot(df['ma_volume_10'], ax=ax2,  width=4,color='#800080', legend='Volume MA10')
    
    # 简化的筹码分布（以收盘价的标准差表示）
    df['price_std'] = df['close'].rolling(window=20).std()
    df['chip_distribution'] = 1 / (1 + df['price_std'] / df['close'])
    
    # 绘制筹码分布
    fplt.plot(df['chip_distribution'], ax=ax3, color='#8B4513', legend='筹码集中度')
    fplt.set_y_range(0, 1.2, ax=ax3)

    # 保存图表
    # 截图并保存

    filename=f'F:/stock_img/100day/{symbol}.png'
    # 显示图表
    fplt.show(qt_exec=True)
    # 等待图表渲染
    time.sleep(1)
    img = fplt.screenshot(filename, fmt='png')
    print(f"图表已保存至: {filename}")
    img_array=np.array(img)
    print(img_array.shape)





def finplot_to_ndarray(plot_func):
    """
    将finplot图表转换为numpy数组

    参数:
        plot_func: 绘制图表的函数

    返回:
        numpy数组，形状为(高度, 宽度, 3)
    """
    # 绘制图表
    plot_func()
    fplt.show(block=False)

    # 截图并转换为ndarray
    img = fplt.screenshot()
    return np.array(img)

def paint(df):
    # 获取股票数据
    print("正在获取股票数据...")
    #df = get_stock_data()
    
    # 计算技术指标
    print("正在计算技术指标...")
    df = calculate_moving_averages(df)
    df = calculate_ene(df)
    
    # 显示数据信息
    print("\n数据基本信息:")
    print(df[['date', 'open', 'high', 'low', 'close', 'volume']].tail())

    # 将df 里的na替换成 浮点数0
    df = df.fillna(0.0)
    # 将df里所有的数字元素精度保留到小数点后5位 需要判断元素是不是数字
    # df = df.applymap(lambda x: x if isinstance(x, (int, float)) else x)

    # 绘制图表
    print("正在绘制图表...")
    plot_stock_chart(df)


def get_all_a_share_tickers():
    """
    获取上海和深圳交易所所有A股股票代码

    Returns:
    tuple: (上海股票代码列表, 深圳股票代码列表)
    """
    try:
        # 获取上海A股实时行情
        sh_stocks = ak.stock_sh_a_spot_em()
        # 转换为FinRL格式
        sh_tickers = [code for code in sh_stocks['代码']]
        print(f"成功获取上海交易所 {len(sh_tickers)} 只股票代码")
    except Exception as e:
        print(f"获取上海股票代码失败: {e}")
        sh_tickers = []

    try:
        # 等待1秒避免请求过于频繁
        time.sleep(1)

        # 获取深圳A股实时行情
        sz_stocks = ak.stock_sz_a_spot_em()
        # 转换为FinRL格式
        sz_tickers = [code for code in sz_stocks['代码']]
        print(f"成功获取深圳交易所 {len(sz_tickers)} 只股票代码")
    except Exception as e:
        print(f"获取深圳股票代码失败: {e}")
        sz_tickers = []

    return sh_tickers, sz_tickers


def get_stock_data(symbol="600036", period="daily", start_date="20250101", end_date="20250731", adjust="qfq"):
    """
    获取股票数据

    Parameters:
    symbol (str): 股票代码
    period (str): 时间周期 ("daily", "weekly", "monthly")
    start_date (str): 开始日期 (YYYYMMDD)
    end_date (str): 结束日期 (YYYYMMDD)
    adjust (str): 复权方式 ("qfq": 前复权, "hfq": 后复权, "": 不复权)

    Returns:
    pd.DataFrame: 股票数据
    """
    df = ak.stock_zh_a_hist(
        symbol=symbol,
        period=period,
        start_date=start_date,
        end_date=end_date,
        adjust=adjust
    )

    # 将中文列名转换为英文列名
    column_mapping = {
        '日期': 'date',
        '股票代码': 'symbol',
        '开盘': 'open',
        '收盘': 'close',
        '最高': 'high',
        '最低': 'low',
        '成交量': 'volume',
        '成交额': 'amount',
        '振幅': 'amplitude',
        '涨跌幅': 'change_rate',
        '涨跌额': 'change_price',
        '换手率': 'turnover_rate'
    }

    # 重命名列
    df = df.rename(columns=column_mapping)
    # 将date 转为时间

    df['timestamp'] = pd.to_datetime(df['date'])
    # 将timestamp 设置成时间索引
    df = df.set_index('timestamp')
    # df = df.reset_index()
    # df = df.rename(columns={"index": "timestamp"})

    return df


def download_all_stock_data(save_path="F:/stock_data/daily_csv"):
    """
    下载所有A股股票数据并保存到指定目录

    Parameters:
    save_path (str): 数据保存路径
    """
    # 确保保存路径存在
    if not os.path.exists(save_path):
        os.makedirs(save_path)
        print(f"创建目录: {save_path}")

    # 获取所有股票代码
    print("正在获取上海和深圳交易所股票代码...")
    sh_tickers, sz_tickers = get_all_a_share_tickers()

    all_tickers = sh_tickers + sz_tickers
    print(f"总共获取到 {len(all_tickers)} 只股票代码")

    # 遍历所有股票代码并下载数据
    success_count = 0
    fail_count = 0

    for i, ticker in enumerate(all_tickers):
        try:
            print(f"[{i + 1}/{len(all_tickers)}] 正在下载股票 {ticker} 的数据...")

            # 获取股票数据
            df = get_stock_data(symbol=ticker, period="daily",
                                start_date="20140101", end_date="20251231", adjust="qfq")

            if not df.empty:
                # 保存为CSV文件
                file_path = os.path.join(save_path, f"{ticker}.csv")
                df.to_csv(file_path, encoding='utf-8-sig')
                print(f"已保存股票 {ticker} 的数据到 {file_path}")
                success_count += 1
            else:
                print(f"股票 {ticker} 没有数据")
                fail_count += 1

            # 等待一段时间避免请求过于频繁
            time.sleep(2.5)

        except Exception as e:
            print(f"下载股票 {ticker} 数据时出错: {e}")
            fail_count += 1
            continue

    print(f"\n下载完成! 成功: {success_count} 只股票, 失败: {fail_count} 只股票")
    print(f"数据保存在: {save_path}")


def download_missing_data(all_tickers,save_path="F:/stock_data/daily_csv"):
    # 遍历所有股票代码并下载数据
    success_count = 0
    fail_count = 0
    for i, ticker in enumerate(all_tickers):
        try:
            print(f"[{i + 1}/{len(all_tickers)}] 正在下载股票 {ticker} 的数据...")

            # 获取股票数据
            df = get_stock_data(symbol=ticker, period="daily",
                                start_date="20140101", end_date="20250731", adjust="qfq")

            if not df.empty:
                # 保存为CSV文件
                file_path = os.path.join(save_path, f"{ticker}.csv")
                df.to_csv(file_path, encoding='utf-8-sig')
                print(f"已保存股票 {ticker} 的数据到 {file_path}")
                success_count += 1
            else:
                print(f"股票 {ticker} 没有数据")
                fail_count += 1

            # 等待一段时间避免请求过于频繁
            time.sleep(2.5)

        except Exception as e:
            print(f"下载股票 {ticker} 数据时出错: {e}")
            fail_count += 1
            continue
    print(f"\n下载完成! 成功: {success_count} 只股票, 失败: {fail_count} 只股票")
    print(f"数据保存在: {save_path}")


def read_stock_data_by_symbol(stock_folder="F:/stock_data/daily_csv", symbol=None, start_date=None, end_date=None):
    """
    从指定文件夹读取特定股票代码的CSV文件数据

    Parameters:
    stock_folder (str): 股票数据文件夹路径
    symbol (str): 股票代码，如果为None则读取所有股票数据
    start_date (str): 开始日期，格式为YYYY-MM-DD
    end_date (str): 结束日期，格式为YYYY-MM-DD

    Returns:
    pd.DataFrame: 股票数据，按symbol和timestamp排序
    """
    # 检查文件夹是否存在
    if not os.path.exists(stock_folder):
        raise FileNotFoundError(f"股票数据文件夹 {stock_folder} 不存在")

    # 如果指定了股票代码，则只读取该股票文件
    if symbol:
        file_path = os.path.join(stock_folder, f"{symbol}.csv")
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"股票数据文件 {file_path} 不存在")

        # 读取数据
        df = pd.read_csv(file_path)
        # 添加股票代码列（如果不存在）
        if 'tic' not in df.columns:
            df['tic'] = symbol
        if 'timestamp' not in df.columns and 'date' in df.columns:
            df['timestamp'] = pd.to_datetime(df['date'])
    else:
        # 读取所有股票文件
        files = [f for f in os.listdir(stock_folder) if f.endswith('.csv')]
        if not files:
            raise ValueError(f"在 {stock_folder} 中未找到CSV文件")

        dataframes = []
        for file in files:
            file_path = os.path.join(stock_folder, file)
            try:
                temp_df = pd.read_csv(file_path)
                # 提取股票代码（文件名不含扩展名）
                ticker = os.path.splitext(file)[0]
                # 添加股票代码列（如果不存在）
                if 'tic' not in temp_df.columns:
                    temp_df['tic'] = ticker
                if 'timestamp' not in temp_df.columns and 'date' in temp_df.columns:
                    temp_df['timestamp'] = pd.to_datetime(temp_df['date'])
                dataframes.append(temp_df)
            except Exception as e:
                print(f"读取文件 {file} 时出错: {e}")
                continue

        if not dataframes:
            raise ValueError("未能读取任何股票数据文件")

        # 合并所有数据
        df = pd.concat(dataframes, ignore_index=True)

    # 转换timestamp列为datetime类型（如果尚未转换）
    if 'timestamp' in df.columns:
        df['timestamp'] = pd.to_datetime(df['timestamp'])
    else:
        raise ValueError("数据中缺少timestamp列")

    # 根据指定的时间范围过滤数据
    if start_date:
        start_date = pd.to_datetime(start_date)
        df = df[df['timestamp'] >= start_date]

    if end_date:
        end_date = pd.to_datetime(end_date)
        df = df[df['timestamp'] <= end_date]

    # 按股票代码和时间戳排序
    df = df.sort_values(['tic', 'timestamp'], ascending=[True, True])



    return df


def read_all_stock_symbols(stock_folder="F:/stock_data/daily_csv"):
    """
    从指定文件夹读取所有股票代码

    Parameters:
    stock_folder (str): 股票数据文件夹路径

    Returns:
    list: 所有股票代码列表
    """
    # 检查文件夹是否存在
    if not os.path.exists(stock_folder):
        raise FileNotFoundError(f"股票数据文件夹 {stock_folder} 不存在")

    # 获取所有CSV文件名（不含扩展名）
    files = [f for f in os.listdir(stock_folder) if f.endswith('.csv')]
    symbols = [os.path.splitext(file)[0] for file in files]

    # 按股票代码排序
    symbols.sort()

    return symbols


# ... existing code ...

def add_action_column_old(df):
    """
    为股票数据添加action列，根据未来价格走势判断操作

    规则:
    1. 如果未来7个交易日涨幅大于20% 并且ma5均线朝上没有回调，action=1 (买入)
    2. 如果未来1天出现跌幅大于等于5%，或者未来2天出现跌幅大于5%，action=2 (卖出)
    3. 否则action=0 (持有/不操作)

    Parameters:
    df (pd.DataFrame): 股票数据，需要包含close列和ma_price_5列

    Returns:
    pd.DataFrame: 添加了action列的数据
    """
    df = df.copy()

    # 确保数据按时间排序
    df = df.sort_values('timestamp').reset_index(drop=True)

    # 计算5日均线（如果不存在）
    if 'ma_price_5' not in df.columns:
        df = calculate_moving_averages(df, [5])

    # 初始化action列为0
    df['action'] = 0

    # 计算每日涨跌幅
    df['pct_change'] = df['change_pct']

    # 遍历数据，为每个交易日判断action
    for i in range(len(df)):
        # 检查是否可以查看未来7天的数据
        if i + 7 >= len(df):
            continue

        # 规则1：未来7个交易日涨幅大于20% 并且ma5均线朝上没有回调
        # 防止除以0
        if df.loc[i, 'close'] != 0:
            future_7d_return = (df.loc[i + 7, 'close'] / df.loc[i, 'close']) - 1
        else:
            future_7d_return = float('inf') if df.loc[i + 7, 'close'] > 0 else 0

        # 检查ma5是否朝上（当前ma5大于7天前的ma5）
        ma5_up = False
        if i >= 7 and 'ma_price_5' in df.columns:
            # 防止除以0
            if df.loc[i - 7, 'ma_price_5'] != 0:
                ma5_up = df.loc[i, 'ma_price_5'] > df.loc[i - 7, 'ma_price_5']

        # 检查期间是否有回调（价格下跌）
        has_pullback = False
        for j in range(1, 8):
            if i + j < len(df):
                # 防止除以0
                if df.loc[i + j - 1, 'close'] != 0:
                    if df.loc[i + j, 'close'] < df.loc[i + j - 1, 'close']:
                        has_pullback = True
                        break

        if future_7d_return > 0.2 and ma5_up and not has_pullback:
            df.loc[i, 'action'] = 1
            continue

        # 规则2：未来1天跌幅大于等于5%，或者未来2天跌幅大于5%
        # 未来1天跌幅
        if i + 1 < len(df):
            # 防止除以0
            if df.loc[i, 'close'] != 0:
                future_1d_return = (df.loc[i + 1, 'close'] / df.loc[i, 'close']) - 1
            else:
                future_1d_return = float('-inf') if df.loc[i + 1, 'close'] > 0 else 0

            if future_1d_return <= -0.05:
                df.loc[i, 'action'] = 2
                continue

        # 未来2天跌幅
        if i + 2 < len(df):
            # 防止除以0
            if df.loc[i, 'close'] != 0:
                future_2d_return = (df.loc[i + 2, 'close'] / df.loc[i, 'close']) - 1
            else:
                future_2d_return = float('-inf') if df.loc[i + 2, 'close'] > 0 else 0

            if future_2d_return <= -0.05:
                df.loc[i, 'action'] = 2
                continue

    # 删除临时列
    if 'pct_change' in df.columns:
        df.drop(columns=['pct_change'], inplace=True)

    return df


# ... existing code ...

def add_action_column(df):
    """
    为股票数据添加action列，根据未来价格走势判断操作

    规则:
    1. 如果未来3天内涨幅大于9% 且未来7天内涨幅大于20% 且期间回撤幅度不超过5%，action=1 (买入)
    2. 如果未来1天出现跌幅大于等于5%，或者未来2天出现跌幅大于5%，action=2 (卖出)
    3. 否则action=0 (持有/不操作)

    Parameters:
    df (pd.DataFrame): 股票数据，需要包含close列和ma_price_5列

    Returns:
    pd.DataFrame: 添加了action列的数据
    """
    df = df.copy()

    # 确保数据按时间排序
    df = df.sort_values('timestamp').reset_index()

    # 计算5日均线（如果不存在）
    if 'ma_price_5' not in df.columns:
        df = calculate_moving_averages(df, [5])

    # 初始化action列为0
    df['action'] = 0

    # 计算每日涨跌幅
    df['pct_change'] = df['change_pct']

    # 用于记录买入信号的列表
    buy_signals = []

    # 遍历数据，为每个交易日判断action
    for i in range(len(df)):
        # 检查是否可以查看未来7天的数据
        if i + 7 >= len(df):
            continue

        # 规则1：未来3天内涨幅大于9% 且未来7天内涨幅大于20% 且期间回撤幅度不超过5%
        future_3d_return = 0
        future_7d_return = 0
        max_drawdown = 0

        # 防止除以0
        if df.loc[i, 'close'] != 0:
            # 计算未来3天涨幅
            future_3d_return = (df.loc[i + 3, 'close'] / df.loc[i, 'close']) - 1

            # 计算未来7天涨幅
            future_7d_return = (df.loc[i + 7, 'close'] / df.loc[i, 'close']) - 1

            # 计算期间最大回撤
            peak_value = df.loc[i, 'close']
            for j in range(1, 8):  # 检查未来7天内的回撤
                if i + j < len(df):
                    current_value = df.loc[i + j, 'close']
                    if current_value > peak_value:
                        peak_value = current_value
                    else:
                        drawdown = (peak_value - current_value) / peak_value if peak_value != 0 else 0
                        max_drawdown = max(max_drawdown, drawdown)

        # 检查买入条件：未来3天涨幅大于9% 且未来7天涨幅大于20% 且期间回撤不超过5%
        if future_3d_return > 0.09 and future_7d_return > 0.20 and max_drawdown <= 0.05:
            df.loc[i, 'action'] = 1
            # 记录买入信号：股票代码和日期
            buy_signals.append({
                'symbol': df.loc[i, 'tic'] if 'tic' in df.columns else 'unknown',
                'date': df.loc[i, 'timestamp'].strftime('%Y-%m-%d')
            })
            continue

        # 规则2：未来1天跌幅大于等于5%，或者未来2天跌幅大于5%
        # 未来1天跌幅
        if i + 1 < len(df):
            # 防止除以0
            if df.loc[i, 'close'] != 0:
                future_1d_return = (df.loc[i + 1, 'close'] / df.loc[i, 'close']) - 1
            else:
                future_1d_return = float('-inf') if df.loc[i + 1, 'close'] > 0 else 0

            if future_1d_return <= -0.05:
                df.loc[i, 'action'] = 2
                continue

        # 未来2天跌幅
        if i + 2 < len(df):
            # 防止除以0
            if df.loc[i, 'close'] != 0:
                future_2d_return = (df.loc[i + 2, 'close'] / df.loc[i, 'close']) - 1
            else:
                future_2d_return = float('-inf') if df.loc[i + 2, 'close'] > 0 else 0

            if future_2d_return <= -0.05:
                df.loc[i, 'action'] = 2
                continue

    # 如果有买入信号，写入buy.csv文件
    if buy_signals:
        buy_df = pd.DataFrame(buy_signals)
        buy_file = "F:/stock_data/tag/buy.csv"
        # 如果文件已存在，则追加写入，否则创建新文件
        if os.path.exists(buy_file):
            buy_df.to_csv(buy_file, mode='a', header=False, index=False)
        else:
            buy_df.to_csv(buy_file, index=False)

    # 删除临时列
    if 'pct_change' in df.columns:
        df.drop(columns=['pct_change'], inplace=True)

    return df



def process_single_stock(symbol, stock_folder):
    """
    处理单个股票数据，添加action列

    Parameters:
    symbol (str): 股票代码
    stock_folder (str): 股票数据文件夹路径

    Returns:
    tuple: (symbol, success_flag, dataframe or None)
    """
    try:
        # 读取单个股票数据
        file_path = os.path.join(stock_folder, f"{symbol}.csv")
        df = pd.read_csv(file_path)
        df['symbol'] = symbol
        # 添加时间戳列（如果不存在）
        if 'timestamp' not in df.columns and 'date' in df.columns:
            df['timestamp'] = pd.to_datetime(df['date'])
        if 'tic' not in df.columns:
            df['tic'] = df['symbol']
        # 设置时间戳为索引并排序
        df = df.set_index('timestamp')
        # 按时间戳排序
        df = df.sort_values(['timestamp'], ascending=[True])

        # 添加action列
        df_with_action = add_action_column(df)
        # column_mapping = {
        #     'date': 'date',
        #     'symbol': 'symbol',
        #     'open': 'open',
        #     'close': 'close',
        #     'high': 'high',
        #     'low': 'low',
        #     'volume': 'volume',
        #     'amount': 'amount',
        #     'amplitude': 'amplitude',
        #     'change_pct': 'change_pct',
        #     '涨跌额': 'change_price',
        #     'action': 'action',
        #     'turnover_rate': 'turnover_rate'
        # }
        # df_with_action = df_with_action.rename(columns=column_mapping)
        df_with_action['symbol'] =symbol
        save_column:list[str]=['date', 'symbol', 'open', 'close', 'high', 'low', 'volume', 'amount', 'amplitude', 'change_pct', 'change_price', 'action', 'turnover_rate']
        # 只保存 需要的列
        df_with_action = df_with_action[save_column]
        # 保存回文件
        df_with_action.to_csv(file_path, index=False)

        print(f"成功为股票 {symbol} 添加action列")
        return symbol, True, df_with_action

    except Exception as e:
        print(f"处理股票 {symbol} 时出错: {e}")
        return symbol, False, None

def add_action_column_for_all_stocks(stock_folder="F:/stock_data/daily_csv"):
    """
    为所有股票数据添加action列

    Parameters:
    stock_folder (str): 股票数据文件夹路径

    Returns:
    dict: 以股票代码为键，添加了action列的数据为值的字典
    """
    # 检查文件夹是否存在
    if not os.path.exists(stock_folder):
        raise FileNotFoundError(f"股票数据文件夹 {stock_folder} 不存在")

    # 获取所有股票代码
    symbols = read_all_stock_symbols(stock_folder)

    result = {}
    success_count = 0
    fail_count = 0

    # 使用线程池并行处理
    with ThreadPoolExecutor(max_workers=40) as executor:
        # 提交所有任务
        future_to_symbol = {
            executor.submit(process_single_stock, symbol, stock_folder): symbol
            for symbol in symbols
        }

        # 处理完成的任务
        for future in as_completed(future_to_symbol):
            symbol, success, df = future.result()
            if success:
                result[symbol] = df
                success_count += 1
            else:
                fail_count += 1

    print(f"处理完成! 成功: {success_count} 只股票, 失败: {fail_count} 只股票")
    return result


def calculate_rsi(df, window=14):
    """
    计算RSI（相对强弱指数）指标

    Parameters:
    df (pd.DataFrame): 股票数据，需要包含'close'列
    window (int): 计算RSI的周期窗口，默认为14

    Returns:
    pd.Series: RSI值序列
    """
    df = df.copy()

    # 计算价格变化
    delta = df['close'].diff()

    # 分离上涨和下跌
    gain = delta.where(delta > 0, 0)
    loss = -delta.where(delta < 0, 0)

    # 计算平均上涨和下跌（使用指数移动平均）
    avg_gain = gain.rolling(window=window, min_periods=1).mean()
    avg_loss = loss.rolling(window=window, min_periods=1).mean()

    # 计算RS
    rs = avg_gain / avg_loss

    # 计算RSI
    rsi = 100 - (100 / (1 + rs))

    return rsi

def calculate_kdj(df, n=9, m1=3, m2=3):
    """
    计算KDJ指标

    Parameters:
    df (pd.DataFrame): 股票数据
    n (int): 周期
    m1 (int): K值平滑周期
    m2 (int): D值平滑周期

    Returns:
    tuple: (K, D, J)
    """
    df = df.copy()

    # 计算未成熟随机值RSV
    low_min = df['low'].rolling(window=n).min()
    high_max = df['high'].rolling(window=n).max()

    # 防止除以0
    rsv = np.where(high_max - low_min != 0,
                   (df['close'] - low_min) / (high_max - low_min) * 100,
                   50)

    # 计算K值
    k = pd.Series(rsv).ewm(alpha=1 / m1, adjust=False).mean()

    # 计算D值
    d = k.ewm(alpha=1 / m2, adjust=False).mean()

    # 计算J值
    j = 3 * k - 2 * d

    return k.values, d.values, j.values

def calculate_ma(df):
    df['ma5'] = df['close'].rolling(window=5).mean()
    df['ma13'] = df['close'].rolling(window=13).mean()
    df['ma20'] = df['close'].rolling(window=20).mean()
    df['ma60'] = df['close'].rolling(window=60).mean()
    df['ma150'] = df['close'].rolling(window=150).mean()

def calculate_macd(df, short=12, long=26, signal=9):
    """
    计算MACD指标

    Parameters:
    df (pd.DataFrame): 股票数据
    short (int): 短期EMA周期
    long (int): 长期EMA周期
    signal (int): 信号线周期

    Returns:
    tuple: (MACD, Signal, Histogram)
    """
    df = df.copy()

    # 计算EMA
    ema_short = df['close'].ewm(span=short, adjust=False).mean()
    ema_long = df['close'].ewm(span=long, adjust=False).mean()

    # 计算MACD线
    macd = ema_short - ema_long

    # 计算信号线
    signal_line = macd.ewm(span=signal, adjust=False).mean()

    # 计算柱状图
    histogram = macd - signal_line

    return macd.values, signal_line.values, histogram.values

# 修改后的calculate_cyq_distribution函数
def calculate_cyq_distribution(df, target_index, days=120, precision=100):
    """
    计算筹码分布(CYQ)指标
    """
    # 确定计算范围
    start_idx = max(0, target_index - days + 1)
    end_idx = target_index + 1
    data_slice = df.iloc[start_idx:end_idx].copy()

    if len(data_slice) == 0:
        return None

    # 计算价格区间
    min_price = data_slice['low'].min()
    max_price = data_slice['high'].max()

    # 防止价格区间为0
    if max_price == min_price:
        max_price = min_price * 1.1 if min_price > 0 else 1.0

    # 使用更精确的价格步长计算
    price_step = (max_price - min_price) / (precision - 1)

    # 初始化筹码分布数组
    cyq_distribution = np.zeros(precision)

    # 计算每天的筹码分布
    for i in range(len(data_slice)):
        row = data_slice.iloc[i]
        open_price, high_price, low_price, close_price = row['open'], row['high'], row['low'], row['close']
        volume = row['volume']

        # 如果没有成交量，跳过
        if volume <= 0:
            continue

        # 计算换手率，如果没有换手率数据则使用默认值
        turnover_rate = row.get('turnover_rate', 100.0) / 100.0 if 'turnover_rate' in row else min(1.0, 100.0 / 100.0)
        turnover_rate = max(0.0, min(1.0, turnover_rate))  # 确保换手率在0-1之间

        # 计算平均价格
        avg_price = (open_price + high_price + low_price + close_price) / 4

        # 调整分布：之前的部分按换手率衰减
        cyq_distribution = cyq_distribution * (1 - turnover_rate)

        # 计算价格分布权重
        # 简化模型：假设价格在高低价格之间呈三角分布，峰值在avg_price
        low_idx = int((low_price - min_price) / price_step)
        high_idx = int((high_price - min_price) / price_step)
        avg_idx = int((avg_price - min_price) / price_step)

        # 确保索引在有效范围内
        low_idx = max(0, min(precision - 1, low_idx))
        high_idx = max(0, min(precision - 1, high_idx))
        avg_idx = max(0, min(precision - 1, avg_idx))

        # 添加当日的筹码分布
        if high_idx == low_idx:
            # 特殊情况：价格区间为0（一字板）
            cyq_distribution[low_idx] += volume * turnover_rate
        else:
            # 正常情况：按三角分布计算
            total_weight = 0
            weights = np.zeros(high_idx - low_idx + 1)

            for j in range(low_idx, high_idx + 1):
                price_at_idx = min_price + j * price_step
                if price_at_idx <= avg_price:
                    # 上半部分三角形
                    if avg_price != low_price:
                        weight = (price_at_idx - low_price) / (avg_price - low_price)
                    else:
                        weight = 1.0
                else:
                    # 下半部分三角形
                    if high_price != avg_price:
                        weight = (high_price - price_at_idx) / (high_price - avg_price)
                    else:
                        weight = 1.0
                weights[j - low_idx] = weight
                total_weight += weight

            # 归一化权重并分配成交量
            if total_weight > 0:
                weights = weights / total_weight
                for j in range(len(weights)):
                    cyq_distribution[low_idx + j] += volume * turnover_rate * weights[j]

    # 归一化分布
    total_volume = np.sum(cyq_distribution)
    if total_volume > 0:
        cyq_distribution = cyq_distribution / total_volume

    # 计算价格轴
    price_axis = np.array([min_price + i * price_step for i in range(precision)])

    # 计算平均成本
    if total_volume > 0:
        avg_cost = np.sum(cyq_distribution * price_axis) / np.sum(cyq_distribution)
    else:
        avg_cost = data_slice.iloc[-1]['close']

    # 计算获利比例（当前价格以下的筹码占比）
    current_price = data_slice.iloc[-1]['close']
    profit_idx = int((current_price - min_price) / price_step)
    profit_ratio = np.sum(cyq_distribution[:profit_idx + 1]) / np.sum(cyq_distribution) if np.sum(
        cyq_distribution) > 0 else 0

    return {
        'distribution': cyq_distribution,
        'price_axis': price_axis,
        'avg_cost': avg_cost,
        'profit_ratio': profit_ratio,
        'current_price': current_price
    }


def calculate_cyq_distribution_v2(df, target_index, days=120, precision=100):
    """
    重新实现的筹码分布函数，更贴近实际市场情况

    Parameters:
    df (pd.DataFrame): 股票数据
    target_index (int): 目标日期在数据中的索引位置
    days (int): 计算周期（天数）
    precision (int): 价格精度（分段数量）

    Returns:
    dict: 包含筹码分布结果的字典
    """
    # 确定计算范围
    start_idx = max(0, target_index - days + 1)
    end_idx = target_index + 1
    data_slice = df.iloc[start_idx:end_idx].copy()

    if len(data_slice) == 0:
        return None

    # 计算价格区间
    min_price = data_slice['low'].min()
    max_price = data_slice['high'].max()

    # 防止价格区间为0
    if max_price == min_price:
        max_price = min_price * 1.1 if min_price > 0 else 1.0

    # 使用更精确的价格步长计算
    price_step = (max_price - min_price) / (precision - 1)

    # 初始化筹码分布数组
    cyq_distribution = np.zeros(precision)

    # 计算每天的筹码分布
    for i in range(len(data_slice)):
        row = data_slice.iloc[i]
        open_price, high_price, low_price, close_price = row['open'], row['high'], row['low'], row['close']
        volume = row['volume']

        # 如果没有成交量，跳过
        if volume <= 0:
            continue

        # 计算换手率，如果没有换手率数据则使用默认值
        turnover_rate = row.get('turnover_rate', 100.0) / 100.0 if 'turnover_rate' in row else min(1.0, 100.0 / 100.0)
        turnover_rate = max(0.0, min(1.0, turnover_rate))  # 确保换手率在0-1之间

        # 计算平均价格
        avg_price = (open_price + high_price + low_price + close_price) / 4

        # 调整分布：之前的部分按换手率衰减
        cyq_distribution = cyq_distribution * (1 - turnover_rate)

        # 计算价格分布权重
        # 改进模型：使用更真实的交易行为模拟
        low_idx = int((low_price - min_price) / price_step)
        high_idx = int((high_price - min_price) / price_step)
        avg_idx = int((avg_price - min_price) / price_step)

        # 确保索引在有效范围内
        low_idx = max(0, min(precision - 1, low_idx))
        high_idx = max(0, min(precision - 1, high_idx))
        avg_idx = max(0, min(precision - 1, avg_idx))

        # 添加当日的筹码分布
        if high_idx == low_idx:
            # 特殊情况：价格区间为0（一字板）
            cyq_distribution[low_idx] += volume * turnover_rate
        else:
            # 正常情况：按更真实的分布计算
            total_weight = 0
            weights = np.zeros(high_idx - low_idx + 1)

            for j in range(low_idx, high_idx + 1):
                price_at_idx = min_price + j * price_step

                # 使用正态分布模拟价格分布
                if price_at_idx <= avg_price:
                    # 上半部分：从低到高逐渐增加
                    weight = (price_at_idx - low_price) / (avg_price - low_price)
                else:
                    # 下半部分：从高到低逐渐减少
                    weight = (high_price - price_at_idx) / (high_price - avg_price)

                weights[j - low_idx] = weight
                total_weight += weight

            # 归一化权重并分配成交量
            if total_weight > 0:
                weights = weights / total_weight
                for j in range(len(weights)):
                    cyq_distribution[low_idx + j] += volume * turnover_rate * weights[j]

    # 归一化分布
    total_volume = np.sum(cyq_distribution)
    if total_volume > 0:
        cyq_distribution = cyq_distribution / total_volume

    # 计算价格轴
    price_axis = np.array([min_price + i * price_step for i in range(precision)])

    # 计算平均成本
    if total_volume > 0:
        avg_cost = np.sum(cyq_distribution * price_axis) / np.sum(cyq_distribution)
    else:
        avg_cost = data_slice.iloc[-1]['close']

    # 计算获利比例（当前价格以下的筹码占比）
    current_price = data_slice.iloc[-1]['close']
    profit_idx = int((current_price - min_price) / price_step)
    profit_ratio = np.sum(cyq_distribution[:profit_idx + 1]) / np.sum(cyq_distribution) if np.sum(
        cyq_distribution) > 0 else 0

    return {
        'distribution': cyq_distribution,
        'price_axis': price_axis,
        'avg_cost': avg_cost,
        'profit_ratio': profit_ratio,
        'current_price': current_price
    }
def save_kline(df, symbol, date, path='F:/stock_data/data1'):
    img_dir = os.path.join(path, symbol)
    os.makedirs(img_dir, exist_ok=True)
    df['date'] = pd.to_datetime(df['date'])
    df = df.set_index('date')
    # 根据action参数确定时间范围
    target_date = pd.to_datetime(date)
    # 确保target_date没有时分秒部分
    target_date = target_date.normalize()
    start_time = target_date - pd.Timedelta(days=60)
    end_time = target_date

    # 计算ENE轨道线
    df = calculate_ene(df, N=10, M1=11, M2=9)

    # 计算压力线
    recent_high, recent_low = calculate_pressure_lines(df, period=20)

    # 计算成交量均线
    df['v_ma5'] = df['volume'].rolling(window=5).mean()
    df['v_ma13'] = df['volume'].rolling(window=13).mean()
    df['v_ma60'] = df['volume'].rolling(window=60).mean()

    # 计算RSI
    df['rsi'] = calculate_rsi(df, window=14)
    df['rsi_30'] = 30
    df['rsi_70'] = 70
    # 确保df索引没有时分秒部分
    df.index = df.index.normalize()

    # 现在筛选时间范围内的数据
    df_plot = df[(df.index >= start_time) & (df.index <= end_time)].copy()

    fig = plt.figure(figsize=(16, 16))
    gs = gridspec.GridSpec(5, 2, height_ratios=[16, 2, 2, 2, 2], width_ratios=[5, 1])
    fig.suptitle(f'{symbol} 股票分析图表 - {date}', fontsize=16)

    # 左侧第一行：K线图
    ax1 = plt.subplot(gs[0, 0])

    # 绘制K线图
    df_plot_ohlc = df_plot[['open', 'high', 'low', 'close']]
    up = df_plot_ohlc[df_plot_ohlc.close >= df_plot_ohlc.open]
    down = df_plot_ohlc[df_plot_ohlc.close < df_plot_ohlc.open]

    # 绘制上涨K线（红色）
    ax1.bar(up.index, up.close - up.open, 0.6, bottom=up.open, color='red', alpha=0.7)
    ax1.bar(up.index, up.high - up.close, 0.1, bottom=up.close, color='red', alpha=0.7)
    ax1.bar(up.index, up.low - up.open, 0.1, bottom=up.open, color='red', alpha=0.7)

    # 绘制下跌K线（绿色）
    ax1.bar(down.index, down.close - down.open, 0.6, bottom=down.open, color='green', alpha=0.7)
    ax1.bar(down.index, down.high - down.open, 0.1, bottom=down.open, color='green', alpha=0.7)
    ax1.bar(down.index, down.low - down.close, 0.1, bottom=down.close, color='green', alpha=0.7)

    # 绘制均线
    ax1.plot(df_plot.index, df_plot['ma5'], label='MA5', color='blue', linewidth=2)
    ax1.plot(df_plot.index, df_plot['ma13'], label='MA13', color='orange', linewidth=1.8)
    ax1.plot(df_plot.index, df_plot['ma20'], label='MA20', color='purple', linewidth=1.6)
    ax1.plot(df_plot.index, df_plot['ma60'], label='MA60', color='brown', linewidth=1.4)
    ax1.plot(df_plot.index, df_plot['ma150'], label='MA150', color='pink', linewidth=1.2)

    # 绘制ENE轨道线
    ax1.plot(df_plot.index, df_plot['ENE_UPPER'], label='ENE_UPPER', color='red', linewidth=1, linestyle='--')
    ax1.plot(df_plot.index, df_plot['ENE_LOWER'], label='ENE_LOWER', color='green', linewidth=1, linestyle='--')
    ax1.plot(df_plot.index, df_plot['ENE'], label='ENE', color='black', linewidth=1)

    # 绘制压力线
    ax1.axhline(y=recent_high, color='red', linestyle='--', alpha=0.7, label=f'压力线: {recent_high:.2f}')
    ax1.axhline(y=recent_low, color='green', linestyle='--', alpha=0.7, label=f'支撑线: {recent_low:.2f}')

    ax1.set_title('K线图与技术指标')
    ax1.set_ylabel('价格')
    ax1.legend(loc='upper left')
    ax1.grid(True, alpha=0.3)

    # 左侧第二行：成交量
    ax2 = plt.subplot(gs[1, 0], sharex=ax1)

    # 绘制成交量柱状图
    ax2.bar(up.index, up['volume'] if 'volume' in up.columns else df_plot.loc[up.index, 'volume'],
            color='red', alpha=0.7, width=0.6)
    ax2.bar(down.index, down['volume'] if 'volume' in down.columns else df_plot.loc[down.index, 'volume'],
            color='green', alpha=0.7, width=0.6)

    # 绘制成交量均线
    ax2.plot(df_plot.index, df_plot['v_ma5'], label='V_MA5', color='blue', linewidth=1.5)
    ax2.plot(df_plot.index, df_plot['v_ma13'], label='V_MA13', color='orange', linewidth=1.2)
    ax2.plot(df_plot.index, df_plot['v_ma60'], label='V_MA60', color='purple', linewidth=1)

    ax2.set_title('成交量')
    ax2.set_ylabel('成交量')
    ax2.legend(loc='upper left')
    ax2.grid(True, alpha=0.3)

    # 左侧第三行：KDJ
    ax3 = plt.subplot(gs[2, 0], sharex=ax1)
    ax3.plot(df_plot.index, df_plot['kdj_k'], label='K', color='blue')
    ax3.plot(df_plot.index, df_plot['kdj_d'], label='D', color='orange')
    ax3.plot(df_plot.index, df_plot['kdj_j'], label='J', color='purple')
    ax3.set_title('KDJ指标')
    ax3.set_ylabel('数值')
    ax3.legend(loc='upper left')
    ax3.grid(True, alpha=0.3)

    # 左侧第四行：RSI
    ax6 = plt.subplot(gs[3, 0], sharex=ax1)
    ax6.plot(df_plot.index, df_plot['rsi'], label='RSI', color='blue')
    ax6.plot(df_plot.index, df_plot['rsi_30'], label='超卖线(30)', color='green', linestyle='--')
    ax6.plot(df_plot.index, df_plot['rsi_70'], label='超买线(70)', color='red', linestyle='--')
    ax6.set_title('RSI指标')
    ax6.set_ylabel('数值')
    ax6.set_ylim(0, 100)
    ax6.legend(loc='upper left')
    ax6.grid(True, alpha=0.3)

    # 左侧第五行：MACD
    ax4 = plt.subplot(gs[4, 0], sharex=ax1)
    # 绘制MACD线和信号线
    ax4.plot(df_plot.index, df_plot['macd'], label='MACD', color='blue')
    ax4.plot(df_plot.index, df_plot['macd_signal'], label='Signal', color='orange')

    # 绘制柱状图
    histogram_positive = np.where(df_plot['macd_histogram'] >= 0, df_plot['macd_histogram'], 0)
    histogram_negative = np.where(df_plot['macd_histogram'] < 0, df_plot['macd_histogram'], 0)
    ax4.bar(df_plot.index, histogram_positive, color='red', alpha=0.7, width=0.6)
    ax4.bar(df_plot.index, histogram_negative, color='green', alpha=0.7, width=0.6)

    ax4.set_title('MACD指标')
    ax4.set_ylabel('数值')
    ax4.legend(loc='upper left')
    ax4.grid(True, alpha=0.3)

    # 右侧：筹码分布图
    ax5 = plt.subplot(gs[0, 1])

    # 设置K线图Y轴范围为最低价*0.9和最高价*1.1
    plot_low = df_plot['low'].min() * 0.9
    plot_high = df_plot['high'].max() * 1.1
    ax1.set_ylim(plot_low, plot_high)

    # 使用CYQ算法计算筹码分布
    # 右侧：筹码分布图
    ax5 = plt.subplot(gs[0, 1])

    # 使用CYQ算法计算筹码分布

    try:
        df_plot_full = df[(df.index <= end_time)].copy()
        # 找到目标日期在数据中的位置
        if target_date in df_plot_full.index:
            target_index = df_plot_full.index.get_loc(target_date)
        else:
            target_index = df_plot_full.index.get_indexer([target_date], method='nearest')[0]

        cyq_df = df_plot_full.reset_index()[['date', 'open', 'close', 'high', 'low', 'volume', 'turnover_rate']].copy()
        if 'turnover_rate' not in cyq_df.columns:
            cyq_df['turnover_rate'] = 100

        cyq_result = calculate_cyq_distribution(cyq_df, target_index, days=120, precision=100)

        if cyq_result is not None and np.sum(cyq_result['distribution']) > 0:
            price_values = cyq_result['price_axis']
            distribution_values = cyq_result['distribution']

            # 放大分布值以增强可视化效果
            max_distribution = np.max(distribution_values)
            if max_distribution > 0:
                distribution_values = distribution_values / max_distribution * (
                        np.max(price_values) - np.min(price_values)) * 0.3

            # 绘制筹码分布曲线（横向）
            ax5.plot(distribution_values, price_values, color='blue', linewidth=1)
            ax5.fill_between(distribution_values, price_values, 0, color='lightblue', alpha=0.7)

            # 标记当前价格和平均成本
            current_price = cyq_result['current_price']
            avg_cost = cyq_result['avg_cost']

            ax5.axhline(y=current_price, color='red', linestyle='--', alpha=0.7, label=f'当前价格: {current_price:.2f}')
            ax5.axhline(y=avg_cost, color='green', linestyle='--', alpha=0.7, label=f'平均成本: {avg_cost:.2f}')

            ax5.set_title('筹码分布')
            ax5.set_xlabel('筹码集中度')
            ax5.set_ylabel('价格')

            # 🔥 关键：同步 Y 轴范围到 K 线图的 Y 轴范围
            # 获取 K 线图的 Y 轴范围
            y_min, y_max = ax1.get_ylim()
            ax5.set_ylim(y_min, y_max)  # 强制设置与 K 线图相同的 Y 轴范围

            ax5.legend(loc='upper left')
            ax5.grid(True, alpha=0.3)
            # 设置筹码分布图Y轴范围与K线图一致
            ax5.set_ylim(plot_low, plot_high)
        else:
            ax5.text(0.5, 0.5, '无有效数据',
                     horizontalalignment='center', verticalalignment='center',
                     transform=ax5.transAxes, fontsize=12)
            ax5.set_title('筹码分布')

    except Exception as e:
        ax5.text(0.5, 0.5, f'筹码分布计算失败\n{str(e)}',
                 horizontalalignment='center', verticalalignment='center',
                 transform=ax5.transAxes, fontsize=12)
        ax5.set_title('筹码分布')

    # 格式化x轴日期显示
    for ax in [ax1, ax2, ax3, ax6, ax4]:
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
        ax.xaxis.set_major_locator(mdates.WeekdayLocator(interval=2))
        plt.setp(ax.xaxis.get_majorticklabels(), rotation=45, ha='right')

    plt.tight_layout()
    savefig_filename = f'{img_dir}/{date.strftime("%Y-%m-%d")}.png'
    fig = plt.gcf()
    fig.savefig(savefig_filename, dpi=72)  # 设置分辨率为300dpi
    plt.close(fig)  # 关闭图像以释放内存



def plot_stock_chart_with_actions(symbol, date, action=True,retrueImage=True):
    """
    使用mplfinance绘制股票图表，包含K线、均线、ENE轨道线、压力线、成交量、KDJ、MACD、RSI和筹码分布

    Parameters:
    symbol (str): 股票代码
    date (str): 指定日期，格式为YYYY-MM-DD
    df (pd.DataFrame): 日线数据
    action (bool): 是否显示动作标记
    retrueImage: 是否返回img 不画图
    """
    df = read_stock_data_by_symbol(symbol=symbol)
    # 确保日期列是datetime类型
    if 'timestamp' in df.columns:
        df['date'] = df['timestamp']
    df['date'] = pd.to_datetime(df['date'])
    df = df.set_index('date')

    # 根据action参数确定时间范围
    target_date = pd.to_datetime(date)
    # 确保target_date没有时分秒部分
    target_date = target_date.normalize()

    if action:
        start_time = target_date - pd.Timedelta(days=60)
        end_time = target_date + pd.Timedelta(days=10)
    else:
        start_time = target_date - pd.Timedelta(days=60)
        end_time = target_date

    if df.empty:
        print("指定日期范围内没有数据")
        return

    # 确保必要的列存在
    required_columns = ['open', 'high', 'low', 'close', 'volume']
    for col in required_columns:
        if col not in df.columns:
            raise ValueError(f"缺少必要的列: {col}")

    # 计算技术指标
    # 计算均线
    df['ma5'] = df['close'].rolling(window=5).mean()
    df['ma13'] = df['close'].rolling(window=13).mean()
    df['ma20'] = df['close'].rolling(window=20).mean()
    df['ma60'] = df['close'].rolling(window=60).mean()
    df['ma150'] = df['close'].rolling(window=150).mean()

    # 计算ENE轨道线
    df = calculate_ene(df, N=10, M1=11, M2=9)

    # 计算压力线
    recent_high, recent_low = calculate_pressure_lines(df, period=20)

    # 计算成交量均线
    df['v_ma5'] = df['volume'].rolling(window=5).mean()
    df['v_ma13'] = df['volume'].rolling(window=13).mean()
    df['v_ma60'] = df['volume'].rolling(window=60).mean()

    # 计算KDJ
    k, d, j = calculate_kdj(df)
    df['kdj_k'] = k
    df['kdj_d'] = d
    df['kdj_j'] = j

    # 计算MACD
    macd, signal, histogram = calculate_macd(df)
    df['macd'] = macd
    df['macd_signal'] = signal
    df['macd_histogram'] = histogram

    # 计算RSI
    df['rsi'] = calculate_rsi(df, window=14)
    df['rsi_30'] = 30
    df['rsi_70'] = 70

    # 确保df索引没有时分秒部分
    df.index = df.index.normalize()

    # 现在筛选时间范围内的数据
    df_plot = df[(df.index >= start_time) & (df.index <= end_time)].copy()

    # 创建图形和子图布局 (增加一个子图位置给RSI)
    fig = plt.figure(figsize=(16,16))
    gs = gridspec.GridSpec(5, 2, height_ratios=[16, 2, 2, 2, 2], width_ratios=[5, 1])
    fig.suptitle(f'{symbol} 股票分析图表 - {date}', fontsize=16)

    # 左侧第一行：K线图
    ax1 = plt.subplot(gs[0, 0])

    # 绘制K线图
    df_plot_ohlc = df_plot[['open', 'high', 'low', 'close']]
    up = df_plot_ohlc[df_plot_ohlc.close >= df_plot_ohlc.open]
    down = df_plot_ohlc[df_plot_ohlc.close < df_plot_ohlc.open]

    # 绘制上涨K线（红色）
    ax1.bar(up.index, up.close - up.open, 0.6, bottom=up.open, color='red', alpha=0.7)
    ax1.bar(up.index, up.high - up.close, 0.1, bottom=up.close, color='red', alpha=0.7)
    ax1.bar(up.index, up.low - up.open, 0.1, bottom=up.open, color='red', alpha=0.7)

    # 绘制下跌K线（绿色）
    ax1.bar(down.index, down.close - down.open, 0.6, bottom=down.open, color='green', alpha=0.7)
    ax1.bar(down.index, down.high - down.open, 0.1, bottom=down.open, color='green', alpha=0.7)
    ax1.bar(down.index, down.low - down.close, 0.1, bottom=down.close, color='green', alpha=0.7)

    # 绘制均线
    ax1.plot(df_plot.index, df_plot['ma5'], label='MA5', color='blue', linewidth=2)
    ax1.plot(df_plot.index, df_plot['ma13'], label='MA13', color='orange', linewidth=1.8)
    ax1.plot(df_plot.index, df_plot['ma20'], label='MA20', color='purple', linewidth=1.6)
    ax1.plot(df_plot.index, df_plot['ma60'], label='MA60', color='brown', linewidth=1.4)
    ax1.plot(df_plot.index, df_plot['ma150'], label='MA150', color='pink', linewidth=1.2)

    # 绘制ENE轨道线
    ax1.plot(df_plot.index, df_plot['ENE_UPPER'], label='ENE_UPPER', color='red', linewidth=1, linestyle='--')
    ax1.plot(df_plot.index, df_plot['ENE_LOWER'], label='ENE_LOWER', color='green', linewidth=1, linestyle='--')
    ax1.plot(df_plot.index, df_plot['ENE'], label='ENE', color='black', linewidth=1)

    # 绘制压力线
    ax1.axhline(y=recent_high, color='red', linestyle='--', alpha=0.7, label=f'压力线: {recent_high:.2f}')
    ax1.axhline(y=recent_low, color='green', linestyle='--', alpha=0.7, label=f'支撑线: {recent_low:.2f}')

    # 标记指定日期
    if action:
        if target_date in df_plot.index:
            close_price = df_plot.loc[target_date, 'close']
            ax1.scatter(target_date, close_price, color='blue', s=100, marker='^', label='目标日期')
        else:
            # 如果指定日期不在索引中，找到最接近的日期
            nearest_date = df_plot.index[df_plot.index.get_indexer([target_date], method='nearest')[0]]
            close_price = df_plot.loc[nearest_date, 'close']
            ax1.scatter(nearest_date, close_price, color='blue', s=100, marker='^', label='目标日期(最近交易日)')

    ax1.set_title('K线图与技术指标')
    ax1.set_ylabel('价格')
    ax1.legend(loc='upper left')
    ax1.grid(True, alpha=0.3)

    # 左侧第二行：成交量
    ax2 = plt.subplot(gs[1, 0], sharex=ax1)

    # 绘制成交量柱状图
    ax2.bar(up.index, up['volume'] if 'volume' in up.columns else df_plot.loc[up.index, 'volume'],
            color='red', alpha=0.7, width=0.6)
    ax2.bar(down.index, down['volume'] if 'volume' in down.columns else df_plot.loc[down.index, 'volume'],
            color='green', alpha=0.7, width=0.6)

    # 绘制成交量均线
    ax2.plot(df_plot.index, df_plot['v_ma5'], label='V_MA5', color='blue', linewidth=1.5)
    ax2.plot(df_plot.index, df_plot['v_ma13'], label='V_MA13', color='orange', linewidth=1.2)
    ax2.plot(df_plot.index, df_plot['v_ma60'], label='V_MA60', color='purple', linewidth=1)

    ax2.set_title('成交量')
    ax2.set_ylabel('成交量')
    ax2.legend(loc='upper left')
    ax2.grid(True, alpha=0.3)

    # 左侧第三行：KDJ
    ax3 = plt.subplot(gs[2, 0], sharex=ax1)
    ax3.plot(df_plot.index, df_plot['kdj_k'], label='K', color='blue')
    ax3.plot(df_plot.index, df_plot['kdj_d'], label='D', color='orange')
    ax3.plot(df_plot.index, df_plot['kdj_j'], label='J', color='purple')
    ax3.set_title('KDJ指标')
    ax3.set_ylabel('数值')
    ax3.legend(loc='upper left')
    ax3.grid(True, alpha=0.3)

    # 左侧第四行：RSI
    ax6 = plt.subplot(gs[3, 0], sharex=ax1)
    ax6.plot(df_plot.index, df_plot['rsi'], label='RSI', color='blue')
    ax6.plot(df_plot.index, df_plot['rsi_30'], label='超卖线(30)', color='green', linestyle='--')
    ax6.plot(df_plot.index, df_plot['rsi_70'], label='超买线(70)', color='red', linestyle='--')
    ax6.set_title('RSI指标')
    ax6.set_ylabel('数值')
    ax6.set_ylim(0, 100)
    ax6.legend(loc='upper left')
    ax6.grid(True, alpha=0.3)

    # 左侧第五行：MACD
    ax4 = plt.subplot(gs[4, 0], sharex=ax1)
    # 绘制MACD线和信号线
    ax4.plot(df_plot.index, df_plot['macd'], label='MACD', color='blue')
    ax4.plot(df_plot.index, df_plot['macd_signal'], label='Signal', color='orange')

    # 绘制柱状图
    histogram_positive = np.where(df_plot['macd_histogram'] >= 0, df_plot['macd_histogram'], 0)
    histogram_negative = np.where(df_plot['macd_histogram'] < 0, df_plot['macd_histogram'], 0)
    ax4.bar(df_plot.index, histogram_positive, color='red', alpha=0.7, width=0.6)
    ax4.bar(df_plot.index, histogram_negative, color='green', alpha=0.7, width=0.6)

    ax4.set_title('MACD指标')
    ax4.set_ylabel('数值')
    ax4.legend(loc='upper left')
    ax4.grid(True, alpha=0.3)

    # 右侧：筹码分布图
    ax5 = plt.subplot(gs[0, 1])

    # 设置K线图Y轴范围为最低价*0.9和最高价*1.1
    plot_low = df_plot['low'].min() * 0.9
    plot_high = df_plot['high'].max() * 1.1
    ax1.set_ylim(plot_low, plot_high)

    # 使用CYQ算法计算筹码分布
    # 右侧：筹码分布图
    ax5 = plt.subplot(gs[0, 1])

    # 使用CYQ算法计算筹码分布


    try:
        df_plot_full = df[(df.index <= end_time)].copy()
        # 找到目标日期在数据中的位置
        if target_date in df_plot_full.index:
            target_index = df_plot_full.index.get_loc(target_date)
        else:
            target_index = df_plot_full.index.get_indexer([target_date], method='nearest')[0]

        cyq_df = df_plot_full.reset_index()[['date', 'open', 'close', 'high', 'low', 'volume', 'turnover_rate']].copy()
        if 'turnover_rate' not in cyq_df.columns:
            cyq_df['turnover_rate'] = 100

        cyq_result = calculate_cyq_distribution(cyq_df, target_index, days=120, precision=100)

        if cyq_result is not None and np.sum(cyq_result['distribution']) > 0:
            price_values = cyq_result['price_axis']
            distribution_values = cyq_result['distribution']

            # 放大分布值以增强可视化效果
            max_distribution = np.max(distribution_values)
            if max_distribution > 0:
                distribution_values = distribution_values / max_distribution * (
                        np.max(price_values) - np.min(price_values)) * 0.3

            # 绘制筹码分布曲线（横向）
            ax5.plot(distribution_values, price_values, color='blue', linewidth=1)
            ax5.fill_between(distribution_values, price_values, 0, color='lightblue', alpha=0.7)

            # 标记当前价格和平均成本
            current_price = cyq_result['current_price']
            avg_cost = cyq_result['avg_cost']

            ax5.axhline(y=current_price, color='red', linestyle='--', alpha=0.7, label=f'当前价格: {current_price:.2f}')
            ax5.axhline(y=avg_cost, color='green', linestyle='--', alpha=0.7, label=f'平均成本: {avg_cost:.2f}')

            ax5.set_title('筹码分布')
            ax5.set_xlabel('筹码集中度')
            ax5.set_ylabel('价格')

            # 🔥 关键：同步 Y 轴范围到 K 线图的 Y 轴范围
            # 获取 K 线图的 Y 轴范围
            y_min, y_max = ax1.get_ylim()
            ax5.set_ylim(y_min, y_max)  # 强制设置与 K 线图相同的 Y 轴范围

            ax5.legend(loc='upper left')
            ax5.grid(True, alpha=0.3)
            # 设置筹码分布图Y轴范围与K线图一致
            ax5.set_ylim(plot_low, plot_high)
        else:
            ax5.text(0.5, 0.5, '无有效数据',
                     horizontalalignment='center', verticalalignment='center',
                     transform=ax5.transAxes, fontsize=12)
            ax5.set_title('筹码分布')

    except Exception as e:
        ax5.text(0.5, 0.5, f'筹码分布计算失败\n{str(e)}',
                 horizontalalignment='center', verticalalignment='center',
                 transform=ax5.transAxes, fontsize=12)
        ax5.set_title('筹码分布')

    # 格式化x轴日期显示
    for ax in [ax1, ax2, ax3, ax6, ax4]:
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
        ax.xaxis.set_major_locator(mdates.WeekdayLocator(interval=2))
        plt.setp(ax.xaxis.get_majorticklabels(), rotation=45, ha='right')

    plt.tight_layout()

    if retrueImage:
        # 将图片保存到文件 中
        savefig_filename = f'1.png'
        fig = plt.gcf()
        fig.savefig(savefig_filename)
        # fig.canvas.draw()
        #
        # image_array = np.array(fig.canvas.renderer._renderer)
        # image_array = image_array[:, :, :3]
        # image_array = np.transpose(image_array, (2, 0, 1))  # c w h
        plt.close(fig)  # 关闭图像以释放内存
        fig.savefig(savefig_filename, dpi=72)  # 设置分辨率为300dpi
        # return image_array.astype(np.float32)
    else:
        plt.show()


def find_stocks_with_early_data(stock_folder="F:/stock_data/daily_csv", cutoff_date="2017-01-01"):
    """
    查找第一个日期小于等于指定日期的股票代码

    Parameters:
    stock_folder (str): 股票数据文件夹路径
    cutoff_date (str): 截止日期，格式为YYYY-MM-DD

    Returns:
    list: 第一个日期小于等于截止日期的股票代码列表
    """
    # 检查文件夹是否存在
    if not os.path.exists(stock_folder):
        raise FileNotFoundError(f"股票数据文件夹 {stock_folder} 不存在")

    # 获取所有CSV文件
    csv_files = [f for f in os.listdir(stock_folder) if f.endswith('.csv')]

    # 转换截止日期为datetime对象
    cutoff_date = pd.to_datetime(cutoff_date)

    # 存储符合条件的股票代码
    qualified_stocks = []

    # 遍历每个CSV文件
    for file in csv_files:
        try:
            # file_path = os.path.join(stock_folder, file)
            # # 只读取第一行数据以提高效率
            # df = pd.read_csv(file_path, nrows=1)
            #
            # # 获取股票代码
            symbol = os.path.splitext(file)[0]
            #
            # # 检查日期列
            # if 'date' in df.columns:
            #     first_date = pd.to_datetime(df['date'].iloc[0])
            # elif 'timestamp' in df.columns:
            #     first_date = pd.to_datetime(df['timestamp'].iloc[0])
            # else:
            #     print(f"文件 {file} 缺少日期列，跳过")
            #     continue
            #
            # # 检查第一个日期是否小于等于截止日期
            # if first_date <= cutoff_date:
            qualified_stocks.append(symbol)

        except Exception as e:
            print(f"处理文件 {file} 时出错: {e}")
            continue

    return qualified_stocks


def train_val_test_split(tics, train_ratio=0.6, val_ratio=0.2, test_ratio=0.2, output_folder="F:/stock_data/tics"):
    """
    将股票代码按照指定比例分成训练集-验证集-测试集，并保存到对应的CSV文件

    Parameters:
    tics (list): 股票代码列表
    train_ratio (float): 训练集比例
    val_ratio (float): 验证集比例
    test_ratio (float): 测试集比例
    output_folder (str): 股票数据文件夹路径
    """
    # 检查比例是否有效
    if abs(train_ratio + val_ratio + test_ratio - 1.0) > 1e-6:
        raise ValueError("训练集、验证集和测试集比例之和必须等于1")
    shuffled_tics = tics.copy()
    random.shuffle(shuffled_tics)
    # 计算各集合的大小
    total_count = len(tics)
    train_count = int(total_count * train_ratio)
    val_count = int(total_count * val_ratio)

    # 分割股票代码列表
    train_tics = shuffled_tics[:train_count]
    val_tics = shuffled_tics[train_count:train_count + val_count]
    test_tics = shuffled_tics[train_count + val_count:]

    # 创建输出目录
    train_dir = os.path.join(output_folder, "train")
    val_dir = os.path.join(output_folder, "val")
    test_dir = os.path.join(output_folder, "test")

    os.makedirs(train_dir, exist_ok=True)
    os.makedirs(val_dir, exist_ok=True)
    os.makedirs(test_dir, exist_ok=True)

    # 保存训练集股票代码
    if train_tics:
        train_df = pd.DataFrame(train_tics, columns=['tic'])
        train_df.to_csv(os.path.join(train_dir, "train_tics.csv"), index=False)

        # 复制对应的股票数据文件
        for tic in train_tics:
            src_file = os.path.join(output_folder, f"{tic}.csv")
            dst_file = os.path.join(train_dir, f"{tic}.csv")
            if os.path.exists(src_file):
                pd.read_csv(src_file).to_csv(dst_file, index=False)

    # 保存验证集股票代码
    if val_tics:
        val_df = pd.DataFrame(val_tics, columns=['tic'])
        val_df.to_csv(os.path.join(val_dir, "val_tics.csv"), index=False)

        # 复制对应的股票数据文件
        for tic in val_tics:
            src_file = os.path.join(output_folder, f"{tic}.csv")
            dst_file = os.path.join(val_dir, f"{tic}.csv")
            if os.path.exists(src_file):
                pd.read_csv(src_file).to_csv(dst_file, index=False)

    # 保存测试集股票代码
    if test_tics:
        test_df = pd.DataFrame(test_tics, columns=['tic'])
        test_df.to_csv(os.path.join(test_dir, "test_tics.csv"), index=False)

        # 复制对应的股票数据文件
        for tic in test_tics:
            src_file = os.path.join(output_folder, f"{tic}.csv")
            dst_file = os.path.join(test_dir, f"{tic}.csv")
            if os.path.exists(src_file):
                pd.read_csv(src_file).to_csv(dst_file, index=False)

    print(f"数据集分割完成:")
    print(f"  训练集: {len(train_tics)} 只股票")
    print(f"  验证集: {len(val_tics)} 只股票")
    print(f"  测试集: {len(test_tics)} 只股票")


def main():
    #paint(get_stock_data())
    # pass
    # download_all_stock_data()
    #missTick=['000956','000866','000832','000827','000817','000805','000787','000660','000658','000653','000621','000618','000602','000588','000583','000578','000569','000556','000549','000542','000535','000527','000522','000515','000508','000412','000406','000405','000047','000015','000013','000003']
    # 获取所有股票代码
    # all_symbols = read_all_stock_symbols("F:/stock_data/daily_csv")
    #download_missing_data(missTick)
    # 读取所有股票数据
    # all_data = read_stock_data_by_symbol("F:/stock_data/daily_csv")

    # 读取特定股票数据
    #single_stock_data = read_stock_data_by_symbol("F:/stock_data/daily_csv", symbol="000001")

    # 读取特定时间范围内的数据
    # time_range_data = read_stock_data_by_symbol(
    #     "F:/stock_data/daily_csv",
    #     start_date="2023-01-01",
    #     end_date="2023-12-31"
    # )

    # add_action_column_for_all_stocks()
    # 002536 2025/8/11
    img= plot_stock_chart_with_actions("002536", "2024-10-10",action=False,retrueImage=True)
    tics = find_stocks_with_early_data()
    # 将tics  按照 6-2-2的比例  分成训练集-验证集-测试集 保存到对应的csv
    #train_val_test_split(tics, train_ratio=0.6, val_ratio=0.2, test_ratio=0.2, output_folder="F:/stock_data/tics")

    print(tics)
    a=1+1

def saveTrainData():
    # 找到股票代码
    tics = find_stocks_with_early_data(cutoff_date="2017-01-01")

    # 定义一个全局的df
    df_all_data = pd.read_csv("F:/stock_data/data1/stock_data.csv", dtype={"symbol": str})
    #遍历
    for tic in tics:
        # 打印现在处理的股票代码
        print(f"处理股票代码: {tic}")
        df = read_stock_data_by_symbol(symbol=tic, start_date="2017-01-01")
        # 打印
        if len( df) > 1000 :
            df['date'] = pd.to_datetime(df['date'])
            df = df.set_index('date')

            macd, signal, histogram = calculate_macd(df)
            df['macd'] = macd
            df['macd_signal'] = signal
            df['macd_histogram'] = histogram
            k, d, j = calculate_kdj(df)
            df['kdj_k'] = k
            df['kdj_d'] = d
            df['kdj_j'] = j
            calculate_ma(df)
            rsi = calculate_rsi(df)
            df['rsi'] = rsi
            df = df.reset_index()
            df = df.round(4)
            # 判断哪些行后面还有足够的数据
            valid_indices = df.index[df.index + 10 < len(df)]

            # 使用向量化操作设置值
            df.loc[valid_indices, 'close_1'] = df.loc[valid_indices + 1, 'ma5'].values
            df.loc[valid_indices, 'close_3'] = df.loc[valid_indices + 3, 'ma5'].values
            df.loc[valid_indices, 'close_5'] = df.loc[valid_indices + 5, 'ma5'].values
            df.loc[valid_indices, 'close_7'] = df.loc[valid_indices + 7, 'ma5'].values
            df.loc[valid_indices, 'close_10'] = df.loc[valid_indices + 10, 'ma5'].values
            df['symbol']=df['tic']
            # 获取 df  index 121 到  末尾 -10的数据
            df_temp = df.iloc[151:len(df)-10]
            # 获取df['action'] >0 的数据
            df_data1 = df_temp[df_temp['action'] > 0]
            # 从df里随机获取len(df_data1)个 action=0的数据
            df_data2 = df_temp[df_temp['action'] == 0].sample(n=min(len(df_data1),50))
            df_data = pd.concat([df_data1, df_data2])
            df_data =df_data.sort_values(['date'], ascending=[True])
            # 判断 tic 是否大于000517  如果大于这保存
            if tic > '000878':
                for index, row in df_data.iterrows():

                    save_kline(df,tic,row['date'])
                df_temp = df_data[['symbol','date','action','close_1','close_3','close_5','close_7','close_10']]
                # df_all_data 添加 df_data
                df_all_data = pd.concat([df_all_data, df_temp], ignore_index=True)
                df_all_data.to_csv('F:/stock_data/data1/stock_data.csv', index=False)



if __name__ == "__main__":
    saveTrainData()
